{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use Evolution Strategy to Play BipedalWalker-v3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "D:\\ProgramData\\Anaconda3\\lib\\site-packages\\ale_py\\roms\\__init__.py:94: DeprecationWarning: Automatic importing of atari-py roms won't be supported in future releases of ale-py. Please migrate over to using `ale-import-roms` OR an ALE-supported ROM package. To make this warning disappear you can run `ale-import-roms --import-from-pkg atari_py.atari_roms`.For more information see: https://github.com/mgbellemare/Arcade-Learning-Environment#rom-management\n",
      "  _RESOLVED_ROMS = _resolve_roms()\n"
     ]
    }
   ],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import sys\n",
    "import logging\n",
    "import itertools\n",
    "\n",
    "import numpy as np\n",
    "np.random.seed(0)\n",
    "import gym\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "logging.basicConfig(level=logging.DEBUG,\n",
    "        format='%(asctime)s [%(levelname)s] %(message)s',\n",
    "        stream=sys.stdout, datefmt='%H:%M:%S')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23:30:37 [INFO] env: <BipedalWalker<BipedalWalker-v3>>\n",
      "23:30:37 [INFO] _action_space: None\n",
      "23:30:37 [INFO] _observation_space: None\n",
      "23:30:37 [INFO] _reward_range: None\n",
      "23:30:37 [INFO] _metadata: None\n",
      "23:30:37 [INFO] _max_episode_steps: 1600\n",
      "23:30:37 [INFO] _elapsed_steps: None\n",
      "23:30:37 [INFO] id: BipedalWalker-v3\n",
      "23:30:37 [INFO] entry_point: gym.envs.box2d:BipedalWalker\n",
      "23:30:37 [INFO] reward_threshold: 300\n",
      "23:30:37 [INFO] nondeterministic: False\n",
      "23:30:37 [INFO] max_episode_steps: 1600\n",
      "23:30:37 [INFO] order_enforce: True\n",
      "23:30:37 [INFO] _kwargs: {}\n",
      "23:30:37 [INFO] _env_name: BipedalWalker\n"
     ]
    }
   ],
   "source": [
    "env = gym.make('BipedalWalker-v3')\n",
    "env.seed(0)\n",
    "for key in vars(env):\n",
    "    logging.info('%s: %s', key, vars(env)[key])\n",
    "for key in vars(env.spec):\n",
    "    logging.info('%s: %s', key, vars(env.spec)[key])\n",
    "\n",
    "def clip_reward(reward):\n",
    "    return np.clip(reward, -1., 1.)\n",
    "reward_clipped_env = gym.wrappers.TransformReward(env, clip_reward)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def play_episode(env, agent, max_episode_steps=None, mode=None, render=False):\n",
    "    observation, reward, done = env.reset(), 0., False\n",
    "    agent.reset(mode=mode)\n",
    "    episode_reward, elapsed_steps = 0., 0\n",
    "    while True:\n",
    "        action = agent.step(observation, reward, done)\n",
    "        if render:\n",
    "            env.render()\n",
    "        if done:\n",
    "            break\n",
    "        observation, reward, done, _ = env.step(action)\n",
    "        episode_reward += reward\n",
    "        elapsed_steps += 1\n",
    "        if max_episode_steps and elapsed_steps >= max_episode_steps:\n",
    "            break\n",
    "    agent.close()\n",
    "    return episode_reward, elapsed_steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ESAgent:\n",
    "    def __init__(self, env=None, weights=None, bias=None):\n",
    "        if weights is not None:\n",
    "            self.weights = weights\n",
    "        else:\n",
    "            self.weights = np.zeros((env.observation_space.shape[0],\n",
    "                    env.action_space.shape[0]))\n",
    "        if bias is not None:\n",
    "            self.bias = bias\n",
    "        else:\n",
    "            self.bias = np.zeros(env.action_space.shape[0])\n",
    "\n",
    "    def reset(self, mode=None):\n",
    "        pass\n",
    "\n",
    "    def close(self):\n",
    "        pass\n",
    "\n",
    "    def step(self, observation, _reward, _done):\n",
    "        action = np.matmul(observation, self.weights)\n",
    "        return action\n",
    "\n",
    "    def train(self, env, scale=0.05, learning_rate=0.2, population=16):\n",
    "        # permulate weights\n",
    "        weight_deltas = [scale * np.random.randn(*agent.weights.shape) for _ in\n",
    "                range(population)]\n",
    "        bias_deltas = [scale * np.random.randn(*agent.bias.shape) for _ in\n",
    "                range(population)]\n",
    "\n",
    "        # calculate rewards\n",
    "        agents = [ESAgent(weights=self.weights + weight_delta,\n",
    "                bias=self.bias + bias_delta) for weight_delta, bias_delta in\n",
    "                zip(weight_deltas, bias_deltas)]\n",
    "        rewards = np.array([play_episode(env, agent)[0] for agent in agents])\n",
    "\n",
    "        # standardize the rewards\n",
    "        std = rewards.std()\n",
    "        if np.isclose(std, 0):\n",
    "            coeffs = np.zeros(population)\n",
    "        else:\n",
    "            coeffs = (rewards - rewards.mean()) / std\n",
    "\n",
    "        # update\n",
    "        weight_updates = sum([coeff * weight_delta for coeff, weight_delta in\n",
    "                zip(coeffs, weight_deltas)])\n",
    "        bias_updates = sum([coeff * bias_delta for coeff, bias_delta in\n",
    "                zip(coeffs, bias_deltas)])\n",
    "        self.weights += learning_rate * weight_updates / population\n",
    "        self.bias += learning_rate * bias_updates / population\n",
    "\n",
    "\n",
    "agent = ESAgent(env=env)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23:30:37 [INFO] ==== train & evaluate ====\n",
      "23:30:46 [DEBUG] evaluate generation 0: reward = -92.09, steps = 110\n",
      "23:30:56 [DEBUG] evaluate generation 1: reward = -92.21, steps = 114\n",
      "23:31:05 [DEBUG] evaluate generation 2: reward = -92.06, steps = 110\n",
      "23:31:15 [DEBUG] evaluate generation 3: reward = -92.14, steps = 106\n",
      "23:31:22 [DEBUG] evaluate generation 4: reward = -92.12, steps = 101\n",
      "23:31:36 [DEBUG] evaluate generation 5: reward = -92.16, steps = 104\n",
      "23:31:43 [DEBUG] evaluate generation 6: reward = -92.11, steps = 102\n",
      "23:31:51 [DEBUG] evaluate generation 7: reward = -92.21, steps = 101\n",
      "23:31:58 [DEBUG] evaluate generation 8: reward = -92.42, steps = 99\n",
      "23:32:05 [DEBUG] evaluate generation 9: reward = -92.54, steps = 97\n",
      "23:32:13 [DEBUG] evaluate generation 10: reward = -92.37, steps = 100\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-5-357323968a3c>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      2\u001b[0m \u001b[0mepisode_rewards\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;33m[\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      3\u001b[0m \u001b[1;32mfor\u001b[0m \u001b[0mgeneration\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mitertools\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcount\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 4\u001b[1;33m     \u001b[0magent\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mreward_clipped_env\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      5\u001b[0m     \u001b[0mepisode_reward\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0melapsed_steps\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mplay_episode\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0menv\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0magent\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mrender\u001b[0m\u001b[1;33m=\u001b[0m \u001b[1;36m1\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      6\u001b[0m     \u001b[0mepisode_rewards\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mepisode_reward\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m<ipython-input-4-764903529a3f>\u001b[0m in \u001b[0;36mtrain\u001b[1;34m(self, env, scale, learning_rate, population)\u001b[0m\n\u001b[0;32m     32\u001b[0m                 bias=self.bias + bias_delta) for weight_delta, bias_delta in\n\u001b[0;32m     33\u001b[0m                 zip(weight_deltas, bias_deltas)]\n\u001b[1;32m---> 34\u001b[1;33m         \u001b[0mrewards\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0marray\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mplay_episode\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0menv\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0magent\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m]\u001b[0m \u001b[1;32mfor\u001b[0m \u001b[0magent\u001b[0m \u001b[1;32min\u001b[0m \u001b[0magents\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     35\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     36\u001b[0m         \u001b[1;31m# standardize the rewards\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m<ipython-input-4-764903529a3f>\u001b[0m in \u001b[0;36m<listcomp>\u001b[1;34m(.0)\u001b[0m\n\u001b[0;32m     32\u001b[0m                 bias=self.bias + bias_delta) for weight_delta, bias_delta in\n\u001b[0;32m     33\u001b[0m                 zip(weight_deltas, bias_deltas)]\n\u001b[1;32m---> 34\u001b[1;33m         \u001b[0mrewards\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0marray\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mplay_episode\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0menv\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0magent\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m]\u001b[0m \u001b[1;32mfor\u001b[0m \u001b[0magent\u001b[0m \u001b[1;32min\u001b[0m \u001b[0magents\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     35\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     36\u001b[0m         \u001b[1;31m# standardize the rewards\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m<ipython-input-3-e346020b060d>\u001b[0m in \u001b[0;36mplay_episode\u001b[1;34m(env, agent, max_episode_steps, mode, render)\u001b[0m\n\u001b[0;32m      4\u001b[0m     \u001b[0mepisode_reward\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0melapsed_steps\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;36m0.\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;36m0\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      5\u001b[0m     \u001b[1;32mwhile\u001b[0m \u001b[1;32mTrue\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 6\u001b[1;33m         \u001b[0maction\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0magent\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mobservation\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mreward\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mdone\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      7\u001b[0m         \u001b[1;32mif\u001b[0m \u001b[0mrender\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      8\u001b[0m             \u001b[0menv\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mrender\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m<ipython-input-4-764903529a3f>\u001b[0m in \u001b[0;36mstep\u001b[1;34m(self, observation, _reward, _done)\u001b[0m\n\u001b[0;32m     18\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     19\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0mstep\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mobservation\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0m_reward\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0m_done\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 20\u001b[1;33m         \u001b[0maction\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmatmul\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mobservation\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mweights\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     21\u001b[0m         \u001b[1;32mreturn\u001b[0m \u001b[0maction\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     22\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "logging.info('==== train & evaluate ====')\n",
    "episode_rewards = []\n",
    "for generation in itertools.count():\n",
    "    agent.train(reward_clipped_env)\n",
    "    episode_reward, elapsed_steps = play_episode(env, agent, render= 1)\n",
    "    episode_rewards.append(episode_reward)\n",
    "    logging.debug('evaluate generation %d: reward = %.2f, steps = %d',\n",
    "            generation, episode_reward, elapsed_steps)\n",
    "    if np.mean(episode_rewards[-10:]) > env.spec.reward_threshold:\n",
    "        break\n",
    "plt.plot(episode_rewards)\n",
    "\n",
    "\n",
    "logging.info('==== test ====')\n",
    "episode_rewards = []\n",
    "for episode in range(100):\n",
    "    episode_reward, elapsed_steps = play_episode(env, agent)\n",
    "    episode_rewards.append(episode_reward)\n",
    "    logging.debug('test episode %d: reward = %.2f, steps = %d',\n",
    "            episode, episode_reward, elapsed_steps)\n",
    "logging.info('average episode reward = %.2f ± %.2f',\n",
    "        np.mean(episode_rewards), np.std(episode_rewards))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
